import numpy as np
import pandas as pd
import torch


# predict_utils.py
def add_time_features(df):
    df['hour_sin'] = np.sin(2 * np.pi * df.index.hour / 24)
    df['hour_cos'] = np.cos(2 * np.pi * df.index.hour / 24)
    df['month_sin'] = np.sin(2 * np.pi * df.index.month / 12)
    df['month_cos'] = np.cos(2 * np.pi * df.index.month / 12)
    df['dayofweek_sin'] = np.sin(2 * np.pi * df.index.dayofweek / 7)
    df['dayofweek_cos'] = np.cos(2 * np.pi * df.index.dayofweek / 7)
    return df


def create_sequences(data, seq_length):
    """
    创建输入序列和目标序列（所有特征均为目标）
    Args:
        data: DataFrame，包含所有特征
        seq_length: 序列长度
    Returns:
        X: 输入序列 (num_samples, seq_length, num_features)
        y: 目标序列 (num_samples, num_features)
    """
    X, y = [], []
    for i in range(len(data) - seq_length):
        X_window = data.iloc[i:i + seq_length].to_numpy()  # 显式转换避免warning
        y_value = data.iloc[i + seq_length].to_numpy()

        X.append(X_window)
        y.append(y_value)
    return np.array(X), pd.DataFrame(y, columns=data.columns)


def predict_time_range(model, data_scaler, data, start_time, num_steps, seq_length=672):
    """
    输入时间点和预测步数，返回多特征预测结果
    Args:
        model: 训练好的模型
        data_scaler: 数据标准化器（StandardScaler）
        data: 原始数据 DataFrame（带时间索引）
        start_time: 起始时间点（字符串，如 '2025-01-15 23:30'）
        num_steps: 需要预测的时间步数
        seq_length: 模型训练时的序列长度
    Returns:
        pred_df: DataFrame，形状为 (num_steps, num_features)
    """
    # 初始化输入序列
    historical_data = data.last(f'{seq_length}T')
    if len(historical_data) < seq_length:
        raise ValueError(f"需要至少 {seq_length} 个历史时间步")

    # 预测循环
    predictions = []
    current_seq = data_scaler.transform(historical_data)
    model.eval()

    with torch.no_grad(), torch.cuda.amp.autocast():
        for _ in range(num_steps):
            input_tensor = torch.FloatTensor(current_seq[-seq_length:]).unsqueeze(0)
            pred, _ = model(input_tensor)
            pred_orig = data_scaler.inverse_transform(pred.cpu().numpy())
            predictions.append(pred_orig[0])
            current_seq = np.vstack([current_seq, pred_orig])

    # 构建结果
    return pd.DataFrame(
        predictions,
        index=pd.date_range(start=start_time, periods=num_steps, freq='30min'),
        columns=data.columns
    )
